import numpy as np
import matplotlib.pyplot as plt

def caculate_nms(boxes, nms_thresh = 0.7):
    keep = []
    if len(boxes) == 0:
        return keep
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    area = (x2 - x1) * (y2 - y1)
    scores = boxes[:,-1]
    order = np.argsort(scores)
    print(order)
    while len(order) > 0:
        i = order[0]
        keep.append(i)
        x11 = np.maximum(x1[i], x1[order[1:]])
        y11 = np.maximum(y1[i], y1[order[1:]])
        x22 = np.minimum(x2[i], x2[order[1:]])
        y22 = np.minimum(y2[i], y2[order[1:]])
        inner_area = (x22 - x11) * (y22 - y11)
        iou = inner_area / (area[i] + area[order[1:]] - inner_area)
        inds = np.where(iou < nms_thresh)[0]
        order = order[inds + 1]
    return keep

def plot_bbox(dets, c='k'):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    plt.plot([x1, x2], [y1, y1], c)
    plt.plot([x1, x1], [y1, y2], c)
    plt.plot([x1, x2], [y2, y2], c)
    plt.plot([x2, x2], [y1, y2], c)
    plt.title("NMS")


if __name__ == "__main__":
    # x0,y0,x1,y1,score
    Boxes = np.array([[100,100,50,50,0.91],
                      [220,220,100,100,0.80],
                      [250,250,70,70,0.99],
                      [300,300,150,150,0.70]])
    plt.figure(1)
    ax1 = plt.subplot(1, 2, 1)
    ax2 = plt.subplot(1, 2, 2)

    # before nms
    plt.sca(ax1)
    plot_bbox(Boxes, 'g')

    # after nms
    keep_boxes = caculate_nms(Boxes)
    plt.sca(ax2)
    plot_bbox(Boxes[keep_boxes], 'r')

    # 设置两个子图的坐标范围，防止坐标范围不一致
    ax1.set_xlim([0, 450])
    ax1.set_ylim([0, 450])
    ax2.set_xlim([0, 450])
    ax2.set_ylim([0, 450])
    plt.show()

